"""File utility functions for radio astronomy simulations."""

from astropy.io import fits
import numpy as np
import os


def read_fits_file(file_path):
    """
    Read FITS file.
    
    Args:
        file_path: Path to FITS file
        
    Returns:
        tuple: (data, header) Data array and header information
    """
    with fits.open(file_path) as hdul:
        data = hdul[0].data
        header = hdul[0].header
        return data, header


def write_fits_file(data, file_path, header=None):
    """
    Write FITS file.
    
    Args:
        data: Data to write
        file_path: Path to output file
        header: Optional header information
    """
    if header is not None:
        hdu = fits.PrimaryHDU(data, header=fits.Header(header))
    else:
        hdu = fits.PrimaryHDU(data)
    
    hdu.writeto(file_path, overwrite=True)


def read_npy_file_sequence(folder_path, prefix='real'):
    """
    Read a sequence of .npy files from a folder.
    Files should be named with a prefix followed by an underscore and number, e.g., real_0.npy, red_1.npy
    
    Args:
        folder_path: Path to folder containing .npy files
        prefix: Prefix of files to load ('real', 'red', or 'red2'), defaults to 'real'
        
    Returns:
        tuple: (frequency, data) Frequency array and combined data from all .npy files
               Frequency is np.linspace(980., 1260., 31) - 31 frequency bins
               Data is a 3D array with shape (30, height, width) representing intensity maps
    """
    import re
    
    # Define frequency bins
    frequency = np.linspace(980., 1260., 31)  # 31 frequency bins
    
    # Get list of .npy files with specified prefix
    files = []
    pattern = f'^{prefix}_(\d+)\.npy$'
    
    for f in os.listdir(folder_path):
        if f.endswith('.npy'):
            # Match pattern like "real_0.npy" or "red_20.npy"
            match = re.match(pattern, f)
            if match:
                files.append((f, int(match.group(1))))
    
    # Sort files by the numeric part
    files.sort(key=lambda x: x[1])
    files = [f[0] for f in files]  # Extract just the filenames
    
    if not files:
        raise FileNotFoundError(f"No valid .npy files with prefix '{prefix}' found in {folder_path}")
    
    if len(files) != 30:
        raise ValueError(f"Expected 30 files with prefix '{prefix}', but found {len(files)}")
    
    # Load all files and combine them
    data_sequence = []
    for file in files:
        file_path = os.path.join(folder_path, file)
        data = np.load(file_path)
        data_sequence.append(data)
    
    # Combine along a new axis (typically frequency axis)
    combined_data = np.stack(data_sequence, axis=0)
    
    return frequency, combined_data


def read_filterbank_file(file_path):
    """
    Read filterbank file using blimpy.
    
    Args:
        file_path: Path to filterbank file
        
    Returns:
        Filterbank object
    """
    try:
        from blimpy import Waterfall
        return Waterfall(file_path)
    except ImportError:
        raise ImportError("blimpy package is required for filterbank file operations")


def write_filterbank_file(frequencies, times, intensities, file_path):
    """
    Write filterbank file sorted by time.
    
    Args:
        frequencies: Frequency array
        times: Time array
        intensities: Intensity array (2D: time x frequency)
        file_path: Path to output file
    """
    try:
        from blimpy import Waterfall
        from blimpy.io import sigproc
    except ImportError:
        raise ImportError("blimpy package is required for filterbank file operations")
    
    # Sort data by time
    time_indices = np.argsort(times)
    sorted_times = times[time_indices]
    sorted_intensities = intensities[time_indices]
    
    # Create filterbank file
    # This is a simplified implementation - in practice, you would need to 
    # properly format the data according to the filterbank specification
    # For now, we'll just indicate that this functionality would be implemented
    pass